#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Pinterest图片下载器
根据搜索结果下载图片到本地
"""

import os
import requests
import time
import random
from urllib.parse import urlparse
from pinterest_image_scraper import PinterestImageScraper
import logging

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class PinterestDownloader:
    def __init__(self, download_dir: str = "pinterest_images"):
        self.download_dir = download_dir
        self.scraper = PinterestImageScraper()
        
        # 创建下载目录
        if not os.path.exists(download_dir):
            os.makedirs(download_dir)
            logger.info(f"创建下载目录: {download_dir}")
    
    def download_images(self, query: str, limit: int = 10) -> list:
        """
        搜索并下载图片
        
        Args:
            query: 搜索关键词
            limit: 下载图片数量
            
        Returns:
            下载成功的文件路径列表
        """
        logger.info(f"开始搜索并下载图片: {query}")
        
        # 获取图片URL列表
        image_urls = self.scraper.search_images(query, limit)
        
        if not image_urls:
            logger.warning("未找到图片URL")
            return []
        
        # 创建查询专用目录
        query_dir = os.path.join(self.download_dir, self._sanitize_filename(query))
        if not os.path.exists(query_dir):
            os.makedirs(query_dir)
        
        downloaded_files = []
        
        for i, url in enumerate(image_urls, 1):
            try:
                logger.info(f"正在下载第 {i}/{len(image_urls)} 张图片...")
                
                # 下载图片
                file_path = self._download_single_image(url, query_dir, i)
                
                if file_path:
                    downloaded_files.append(file_path)
                    logger.info(f"下载成功: {os.path.basename(file_path)}")
                else:
                    logger.warning(f"下载失败: {url}")
                
                # 添加延迟避免被限制
                time.sleep(random.uniform(1, 3))
                
            except Exception as e:
                logger.error(f"下载图片时发生错误: {str(e)}")
                continue
        
        logger.info(f"下载完成! 成功下载 {len(downloaded_files)} 张图片到 {query_dir}")
        return downloaded_files
    
    def _download_single_image(self, url: str, save_dir: str, index: int) -> str:
        """
        下载单张图片
        
        Args:
            url: 图片URL
            save_dir: 保存目录
            index: 图片索引
            
        Returns:
            保存的文件路径，失败返回None
        """
        try:
            # 解析URL获取文件扩展名
            parsed_url = urlparse(url)
            file_ext = os.path.splitext(parsed_url.path)[1] or '.jpg'
            
            # 生成文件名
            filename = f"image_{index:03d}{file_ext}"
            file_path = os.path.join(save_dir, filename)
            
            # 如果文件已存在，跳过下载
            if os.path.exists(file_path):
                logger.info(f"文件已存在，跳过: {filename}")
                return file_path
            
            # 下载图片
            headers = {
                'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/139.0.0.0 Safari/537.36',
                'Referer': 'https://www.pinterest.com/',
            }
            
            response = requests.get(url, headers=headers, timeout=30, stream=True)
            response.raise_for_status()
            
            # 保存图片
            with open(file_path, 'wb') as f:
                for chunk in response.iter_content(chunk_size=8192):
                    if chunk:
                        f.write(chunk)
            
            return file_path
            
        except Exception as e:
            logger.error(f"下载单张图片失败 {url}: {str(e)}")
            return None
    
    def _sanitize_filename(self, filename: str) -> str:
        """
        清理文件名，移除不合法字符
        
        Args:
            filename: 原始文件名
            
        Returns:
            清理后的文件名
        """
        # 移除或替换不合法字符
        invalid_chars = '<>:"/\\|?*'
        for char in invalid_chars:
            filename = filename.replace(char, '_')
        
        # 限制长度
        return filename[:50]


def main():
    """
    主函数 - 命令行接口
    """
    import sys
    
    if len(sys.argv) < 2:
        print("使用方法: python pinterest_downloader.py <搜索关键词> [下载数量]")
        print("示例: python pinterest_downloader.py '美食' 5")
        sys.exit(1)
    
    query = sys.argv[1]
    limit = int(sys.argv[2]) if len(sys.argv) > 2 else 10
    
    print(f"🔍 搜索关键词: {query}")
    print(f"📥 下载数量: {limit}")
    print("=" * 50)
    
    downloader = PinterestDownloader()
    downloaded_files = downloader.download_images(query, limit)
    
    if downloaded_files:
        print(f"\n✅ 下载完成! 共下载 {len(downloaded_files)} 张图片")
        print("📁 文件列表:")
        for file_path in downloaded_files:
            print(f"   - {file_path}")
    else:
        print("\n❌ 没有成功下载任何图片")


if __name__ == "__main__":
    main()
